#ifndef OMPL_GEOMETRIC_PLANNERS_APPROX_CSPACE_ORACLE_H
#define OMPL_GEOMETRIC_PLANNERS_APPROX_CSPACE_ORACLE_H

#include <ompl/base/Planner.h>
#include <ompl/contrib/hplanner/StateTransformer.h>


#include <ompl/hash/topk.h>
#include <ompl/hash/dataset.h>
#include <ompl/hash/LSH.h>
#include <ompl/hash/composite.h>
#include <ompl/hash/LSH_index.h>
#include <ompl/hash/hash_table.h>
#include <ompl/hash/bruteforce.h>
#include <ompl/hash/multiprobeLSH.h>
#include <ompl/hash/embedding.h>
#include <ompl/hash/eval.h>
#include <ompl/hash/common.h>

#include <vector>

#define CONSIDER_NONEUCLIDEAN 1


namespace ompl
{
namespace base
{

#define ApproxCSpaceOracle_PI 3.1415926



class ApproxCSpaceOracleBase
{
public:

  ApproxCSpaceOracleBase(const SpaceInformationPtr& si, bool usePositiveState = true) : si_(si), svt_(si)
  {
    usePositiveState_ = usePositiveState;
    dim_ = svt_.dim();
  }

  virtual ~ApproxCSpaceOracleBase() {}

  bool getUsePositiveState() const
  {
    return usePositiveState_;
  }

  bool& getUsePositiveState()
  {
    return usePositiveState_;
  }

  // insert the state into the knn storage and index it immediately
  virtual void insertAndIndex(const State* state, bool collision) = 0;

  // only insert the state into the knn storage, index it later (not now)
  virtual void insert(const State* state, bool collision) = 0;

  // index the i-th element in the dataset
  virtual void index(unsigned int i, bool collision) = 0;

  virtual std::pair<float, float> computeEmbedParamFromDataset() const;

  virtual std::pair<float, float> computeEmbedParamDefault() const;

  virtual void pointKNNQuery(const State* state, bool collision, std::vector<std::pair<unsigned int, float> >& res) const = 0;

  virtual void lineKNNQuery(const State* s1, const State* s2, bool collision, std::vector<std::pair<unsigned int, float> >& res,
                            float* line_origin = NULL, float* line_direction = NULL, float* line_len = NULL) const = 0;

  const SpaceInformationPtr& getSpaceInformation() const
  {
    return si_;
  }

  const StateToVectorTransformer& getStateToVectorTransformer() const
  {
    return svt_;
  }


  State* getNegativeState(unsigned int i) const
  {
    return collision_samples_[i];
  }

  State* getPositiveState(unsigned int i) const
  {
    return collision_free_samples_[i];
  }

  unsigned int getNumNegativeState() const
  {
    return collision_samples_.size();
  }

  unsigned int getNumPositiveState() const
  {
    return collision_free_samples_.size();
  }

  virtual float* getNegativeRawState(unsigned int i) const = 0;

  virtual float* getPositiveRawState(unsigned int i) const = 0;

  virtual unsigned int getNumNegativeRawState() const = 0;

  virtual unsigned int getNumPositiveRawState() const = 0;


protected:

  SpaceInformationPtr si_;
  mutable StateToVectorTransformer svt_;
  unsigned int dim_;

  std::vector<State*> collision_samples_;
  std::vector<State*> collision_free_samples_;

  bool usePositiveState_;
};



class ApproxCSpaceOracleBasic : public ApproxCSpaceOracleBase
{
public:
  typedef typename cspace_learning::Tail<typename cspace_learning::RepeatHash<typename cspace_learning::GaussianLSH> > MyLSH;
  typedef typename cspace_learning::LinearHashTable<unsigned int, unsigned int, __gnu_cxx::hash<unsigned int> > HashTable;
  typedef typename cspace_learning::LSHIndex<MyLSH, HashTable, unsigned int> Index;
  typedef typename cspace_learning::Dataset Dataset;
  typedef typename cspace_learning::Dataset::Accessor Accessor;
  typedef typename cspace_learning::DefaultRng DefaultRng;
  typedef typename cspace_learning::metric::l2sqr<float> L2sqr;
  typedef typename cspace_learning::dotprod<float> Dotprod;
  typedef typename cspace_learning::ReducedEmbedding Embedding;
  typedef typename cspace_learning::eval::PointToLineEval PointToLineEval;
  typedef typename cspace_learning::eval::PointEval<L2sqr> PointEval;
  typedef typename cspace_learning::TopkScanner<Accessor, const float*, PointToLineEval> LineScanner;
  typedef typename cspace_learning::TopkScanner<Accessor, const float*, PointEval> PointScanner;


  ApproxCSpaceOracleBasic(const SpaceInformationPtr& si,
                          unsigned int K1, float R1, unsigned int M1, unsigned int L1, unsigned int H1, float W1,
                          unsigned int K2, float R2, unsigned int M2, unsigned int L2, unsigned int H2, float W2,
                          bool usePositiveState);

  virtual ~ApproxCSpaceOracleBasic();

  void insert(const State* state, bool collision);

  void insertAndIndex(const State* state, bool collision);

  void index(unsigned int id, bool collision);

  void pointKNNQuery(const State* state, bool collision,
                     std::vector<std::pair<unsigned int, float> >& res) const;

  void lineKNNQuery(const State* s1, const State* s2, bool collision,
                    std::vector<std::pair<unsigned int, float> >& res,
                    float* line_origin = NULL, float* line_direction = NULL, float* line_len = NULL) const;

  float* getNegativeRawState(unsigned int i) const
  {
    return (*collision_dataset_)[i];
  }

  float* getPositiveRawState(unsigned int i) const
  {
    if(!getUsePositiveState()) return NULL;
    return (*collision_free_dataset_)[i];
  }

  unsigned int getNumNegativeRawState() const
  {
    return collision_dataset_->size();
  }

  unsigned int getNumPositiveRawState() const
  {
    if(!getUsePositiveState()) return 0;
    return collision_free_dataset_->size();
  }

protected:


  Dataset* collision_dataset_;
  Dataset* collision_free_dataset_;

  Accessor* collision_accessor_;
  Accessor* collision_free_accessor_;


  DefaultRng rng_;

  unsigned int embed_dim_;

  mutable Index collision_line_index_;
  mutable Index collision_free_line_index_;

  mutable Index collision_point_index_;
  mutable Index collision_free_point_index_;

  L2sqr* line_l2_;
  L2sqr* point_l2_;

  Dotprod* l2dot_;
  Embedding* embed_;

  float collision_embed_c_;
  float collision_free_embed_c_;

  PointToLineEval* line_eval_;
  PointEval* point_eval_;

  LineScanner* collision_line_scanner_;
  LineScanner* collision_free_line_scanner_;

  PointScanner* collision_point_scanner_;
  PointScanner* collision_free_point_scanner_;


  float* transformed_state_1_cache_;
  float* transformed_state_2_cache_;

  float* line_direction_cache_;
  mutable float line_len_cache_;
  float* line_origin_cache_;

  float* transform_state_cache_;
};


inline ApproxCSpaceOracleBasic* generateBasicOracleSlow(const SpaceInformationPtr& si, bool usePositiveState = true)
{
  unsigned int K1 = 50, K2 = 50;
  float R1 = std::numeric_limits<float>::max(), R2 = std::numeric_limits<float>::max();

  unsigned int M1 = 5, M2 = 5;
  unsigned int L1 = 20, L2 = 20;
  unsigned int H1 = 1000, H2 = 1000;
  float W1 = 2.0, W2 = 2.0;

  ApproxCSpaceOracleBasic* oracle = new ApproxCSpaceOracleBasic(si,
                                                                K1, R1, M1, L1, H1, W1,
                                                                K2, R2, M2, L2, H2, W2,
                                                                usePositiveState);
  return oracle;

}

inline ApproxCSpaceOracleBasic* generateBasicOracle(const SpaceInformationPtr& si, bool usePositiveState = true)
{
  unsigned int K1 = 50, K2 = 50;
  float R1 = std::numeric_limits<float>::max(), R2 = std::numeric_limits<float>::max();

  unsigned int M1 = 5, M2 = 5;
  unsigned int L1 = 20, L2 = 20;
  unsigned int H1 = 1000, H2 = 100000;
  float W1 = 2.0, W2 = 1.0;

  ApproxCSpaceOracleBasic* oracle = new ApproxCSpaceOracleBasic(si,
                                                                K1, R1, M1, L1, H1, W1,
                                                                K2, R2, M2, L2, H2, W2,
                                                                usePositiveState);
  return oracle;
}



class ApproxCSpaceOracleMultiprobe : public ApproxCSpaceOracleBase
{
public:

  typedef typename cspace_learning::MultiProbeLSH MyLSH;
  typedef typename cspace_learning::LinearHashTable<unsigned int, unsigned int, __gnu_cxx::hash<unsigned int> > HashTable;
  typedef typename cspace_learning::MultiProbeLSHIndex<HashTable, unsigned int> Index;
  typedef typename cspace_learning::Dataset Dataset;
  typedef typename cspace_learning::Dataset::Accessor Accessor;
  typedef typename cspace_learning::DefaultRng DefaultRng;
  typedef typename cspace_learning::metric::l2sqr<float> L2sqr;
  typedef typename cspace_learning::dotprod<float> Dotprod;
  typedef typename cspace_learning::ReducedEmbedding Embedding;
  typedef typename cspace_learning::eval::PointToLineEval PointToLineEval;
  typedef typename cspace_learning::eval::PointEval<L2sqr> PointEval;
  typedef typename cspace_learning::TopkScanner<Accessor, const float*, PointToLineEval> LineScanner;
  typedef typename cspace_learning::TopkScanner<Accessor, const float*, PointEval> PointScanner;


  ApproxCSpaceOracleMultiprobe(const SpaceInformationPtr& si,
                               unsigned int K1, float R1, unsigned int M1, unsigned int L1, unsigned int H1, float W1,
                               unsigned int K2, float R2, unsigned int M2, unsigned int L2, unsigned int H2, float W2,
                               bool usePositiveState);

  virtual ~ApproxCSpaceOracleMultiprobe();

  void insert(const State* state, bool collision);

  void insertAndIndex(const State* state, bool collision);

  void index(unsigned int id, bool collision);

  void pointKNNQuery(const State* state, bool collision,
                     std::vector<std::pair<unsigned int, float> >& res) const;

  void lineKNNQuery(const State* s1, const State* s2, bool collision,
                    std::vector<std::pair<unsigned int, float> >& res,
                    float* line_origin = NULL, float* line_direction = NULL, float* line_len = NULL) const;

  float* getNegativeRawState(unsigned int i) const
  {
    return (*collision_dataset_)[i];
  }

  float* getPositiveRawState(unsigned int i) const
  {
    if(!getUsePositiveState()) return NULL;
    return (*collision_free_dataset_)[i];
  }

  unsigned int getNumNegativeRawState() const
  {
    return collision_dataset_->size();
  }

  unsigned int getNumPositiveRawState() const
  {
    if(!getUsePositiveState()) return 0;
    return collision_free_dataset_->size();
  }

protected:

  Dataset* collision_dataset_;
  Dataset* collision_free_dataset_;

  Accessor* collision_accessor_;
  Accessor* collision_free_accessor_;


  DefaultRng rng_;

  unsigned int embed_dim_;

  mutable Index collision_line_index_;
  mutable Index collision_free_line_index_;

  mutable Index collision_point_index_;
  mutable Index collision_free_point_index_;

  L2sqr* line_l2_;
  L2sqr* point_l2_;

  Dotprod* l2dot_;
  Embedding* embed_;

  float collision_embed_c_;
  float collision_free_embed_c_;

  PointToLineEval* line_eval_;
  PointEval* point_eval_;

  LineScanner* collision_line_scanner_;
  LineScanner* collision_free_line_scanner_;

  PointScanner* collision_point_scanner_;
  PointScanner* collision_free_point_scanner_;


  float* transformed_state_1_cache_;
  float* transformed_state_2_cache_;

  float* line_direction_cache_;
  mutable float line_len_cache_;
  float* line_origin_cache_;

  float* transform_state_cache_;

  unsigned int multiprobe_num_;
};


inline ApproxCSpaceOracleMultiprobe* generateMultiprobeOracleSlow(const SpaceInformationPtr& si, bool usePositiveState = true)
{
  unsigned int K1 = 50, K2 = 50;
  float R1 = std::numeric_limits<float>::max(), R2 = std::numeric_limits<float>::max();

  unsigned int M1 = 5, M2 = 5;
  unsigned int L1 = 20, L2 = 20;
  unsigned int H1 = 1000, H2 = 1000;
  float W1 = 2.0, W2 = 2.0;

  ApproxCSpaceOracleMultiprobe* oracle = new ApproxCSpaceOracleMultiprobe(si,
                                                                          K1, R1, M1, L1, H1, W1,
                                                                          K2, R2, M2, L2, H2, W2,
                                                                          usePositiveState);
  return oracle;

}

inline ApproxCSpaceOracleMultiprobe* generateMultiprobeOracle(const SpaceInformationPtr& si, bool usePositiveState = true)
{
  unsigned int K1 = 50, K2 = 50;
  float R1 = std::numeric_limits<float>::max(), R2 = std::numeric_limits<float>::max();

  unsigned int M1 = 5, M2 = 5;
  unsigned int L1 = 20, L2 = 20;
  unsigned int H1 = 1000, H2 = 100000;
  float W1 = 2.0, W2 = 1.0;

  ApproxCSpaceOracleMultiprobe* oracle = new ApproxCSpaceOracleMultiprobe(si,
                                                                          K1, R1, M1, L1, H1, W1,
                                                                          K2, R2, M2, L2, H2, W2,
                                                                          usePositiveState);
  return oracle;
}



}
}

#endif
